#include "molsys.h"
#include "constants.h"
#include "utility.h"
#include <fstream>
#include <iostream>
#include <sstream>
#include <vector>
#include <math.h>

//
double MolSys::ComputeTemperature()
{
    ComputeKineticEnergy();
    temperature_ = kinetic_energy_ * constants::Unit_au2J * 2.0 / (3.0 * constants::Boltzmann);
    return temperature_; // in K
}

//
double MolSys::ComputeKineticEnergy()
{
    double temp = 0;
    for (size_t i = 0; i < atoms_.size(); i++)
    {
        for (int j = 0; j < 3; j++)
        {
            temp += atoms_[i].mass_ * atoms_[i].veloc_[j] * atoms_[i].veloc_[j];
        }
    }
    kinetic_energy_ = temp * 0.5;
    return kinetic_energy_;
}

//
void MolSys::ReadMolSys(const std::string &filename)
{
    std::ifstream infile(filename, std::ios::in);
    if (!infile)
    {
        std::cerr << "Error: file '" << filename << "' does NOT exist" << std::endl;
        exit(1);
    }
    // we need xyz file to read
    std::string suffix = filename.substr(filename.find_last_of('.') + 1);
    if (suffix != "xyz")
    {
        std::cerr << "Error: *.xyz file is required to set the system" << std::endl;
        exit(1);
    } // @TODO, read from other filetypes, e.g. pdb, mol, etc.
    // now read xyz file and set atoms
    int nLine = 0;
    atoms_.clear();
    std::string line;
    while (std::getline(infile, line))
    {
        if (line[0] == '#') // comment line
        {
            continue;
        }
        if (line == "") // blank line
        {
            continue;
        }
        nLine += 1;
        if (nLine == 1)
        {
            natoms_ = std::stoi(line);
        }
        if (nLine == 2) // second line is comment
        {
            continue;
        }
        std::vector<std::string> fields = Split(line, " ");
        if (fields.size() >= 4 && nLine - 3 < natoms_) // name x y z (coordinates in Angstrom)
        {
            Atom atom;
            atom.SetName(fields[0]);
            double mass = constants::AtomMassMap.at(fields[0]) * constants::Unit_gpermol2au;
            atom.SetMass(mass); // in au
            std::vector<double> coord;
            coord.push_back(std::stod(fields[1]) * constants::Unit_Angstrom2au);
            coord.push_back(std::stod(fields[2]) * constants::Unit_Angstrom2au);
            coord.push_back(std::stod(fields[3]) * constants::Unit_Angstrom2au);
            atom.SetCoord(coord);   // in au
            if (fields.size() == 7) // name x y z vx vy vz (velocites in Angstrom/fs)
            {
                std::vector<double> veloc;
                veloc.push_back(std::stod(fields[1]) * constants::Unit_Angstrom2au / constants::Unit_fs2au);
                veloc.push_back(std::stod(fields[2]) * constants::Unit_Angstrom2au / constants::Unit_fs2au);
                veloc.push_back(std::stod(fields[3]) * constants::Unit_Angstrom2au / constants::Unit_fs2au);
                atom.SetVeloc(veloc); // in au
            }
            atoms_.push_back(atom);
        }
    }
    if (natoms_ != int(atoms_.size()))
    {
        std::cerr << "Error when reading xyz file" << std::endl;
        exit(1);
    }
    return;
}

//
void MolSys::WriteMolSys(const std::string &filename, const int count)
{
    std::ofstream outfile;
    outfile.open(filename, std::ios::app);
    outfile << natoms_ << std::endl;
    outfile << "iStep = " << count
            << ", Ep(au) = " << potential_energy_ 
            << ", Ek(au) = " << kinetic_energy_
            << ", Etot(au) = " << total_energy_ 
            << ", T(K) = " << temperature_
            << std::endl;
    // write coordinates and velocities, units converted from au
    for (size_t i = 0; i < atoms_.size(); i++)
    {
        outfile << atoms_[i].name_;
        for (size_t j = 0; j < 3; j++)
        {
            outfile << " " << atoms_[i].coord_[j] * constants::Unit_au2Angstrom;
        }
        for (size_t j = 0; j < 3; j++)
        {
            outfile << " " << atoms_[i].veloc_[j] * constants::Unit_au2Angstrom / constants::Unit_au2fs;
        }
        outfile << std::endl;
    }
    outfile.close();
}

//
void MolSys::SetAtomsInCubicBox(const std::string &atom_name, const int natoms, const double &box_length)
{
    if (natoms < 1)
    {
        std::cerr << "Error: natoms too small" << std::endl;
        exit(1);
    }
    natoms_ = natoms;
    atoms_.resize(natoms);
    int nInAxis = (int)ceil(cbrt((double)natoms)); // cubic root
    double shift = box_length / (double)nInAxis;
    std::vector<double> origin({0, 0, 0});
    int count = 0;
    for (int i = 0; i < nInAxis; i++)
    {
        for (int j = 0; j < nInAxis; j++)
        {
            for (int k = 0; k < nInAxis; k++)
            {
                if (count == natoms)
                {
                    return;
                }
                atoms_[count].SetName(atom_name);
                atoms_[count].SetMass(constants::AtomMassMap.at(atom_name));
                atoms_[count].SetCoord({origin[0] + i * shift, origin[1] + j * shift, origin[2] + k * shift});
                count += 1;
            }
        }
    }
    return;
}

void MolSys::RemoveCenterVeloc()
{
    std::vector<double> v_center({0, 0, 0});
    for (int i = 0; i < natoms_; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            v_center[j] += atoms_[i].veloc_[j];
        }
    }
    for (int j = 0; j < 3; j++)
    {
        v_center[j] /= (double)natoms_;
    }
    for (int i = 0; i < natoms_; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            atoms_[i].veloc_[j] -= v_center[j];
        }
    }
}
